from openai import OpenAI
import lex
import re

my_model='替换成本地部署的模型名称，如deepseek-r1:1.5b'

global message_table
message_table = {}

global client
client = OpenAI(
    base_url='http://localhost:11434/v1/',
    api_key='ollama',
)

def change_client(new_client, new_model):
    global my_model
    global client
    client = new_client
    my_model = new_model

# 添加用户
def add_user(user_name, prompt = None):
    message_table[user_name] = []

    if prompt is not None:
        add_prompt(user_name, prompt)

# 添加消息到历史记录
def add_history(user_name, role, content):
    # if role == "user":
    #     str_list=re.split(":",content.strip(),1)
    #     if len(str_list) >= 1:
    #         print(f"content: {str_list[1]}, name: {str_list[0]}")
    #         message_table[user_name].append({"role": role, "content": str_list[1], "name": str_list[0]})
    # else:
    message_table[user_name].append({"role": role, "content": content})

# 将消息发送给deepseek
def chat_with_assistant(message):
    try:
        response = client.chat.completions.create(
            model=my_model,
            messages=message
        )
        content = response.choices[0].message.content
        return content
    
    except Exception as e:
        print(f"Error: {e}")
        return "抱歉，我遇到了一个错误，请稍后再试。"

# 聊天
def chat(user, prompts):
    for prompt in prompts:
        add_history(user, "user", prompt)

    content = chat_with_assistant(message_table[user])

    content = lex.parse(content)
    add_history(user, "assistant", content)

    return content

# 输入提示
def add_prompt(user_name, prompt):
    message_table[user_name].append({"role": "system", "content": prompt})